#!/usr/bin/env python3
# Copyright (C) 2025 lelele group. All rights reserved.
# SPDX-License-Identifier: MIT
"""
Verification of campaign script for Kyoto Cabinet benchmarks using different locks.
"""

import subprocess
import os
from pathlib import Path

from benchkit.campaign import Platform, get_current_platform
from benchkit.utils.dir import caller_dir, get_curdir
from tiltlib import TiltLib

this_dir = caller_dir()
tilt_locks_dir = (this_dir / "../tilt/locks").resolve()
vsync_dir = (tilt_locks_dir / "../deps/libvsync/").resolve()
patch_path = (tilt_locks_dir / "prefetch.diff").resolve()


def build_locks(platform: Platform) -> Path:
    """Builds the Vsync lock configurations."""
    tilt_ok = TiltLib(tilt_locks_dir=tilt_locks_dir, build_prefix="build_ok")
    platform.comm.shell(
        command="git checkout -- include/vsync/atomic/internal/arm64.h",
        current_dir=vsync_dir,
    )
    tilt_ok.build()
    build_ok = tilt_ok.build_dir

    return build_ok


def run_benchmark(command: str, kyotocabinet_src_dir: Path, preload_lib: Path = None) -> None:
    """Run the benchmark command in the correct directory and suppress output."""
    try:
        # Print the current working directory
        print(f"[CD] {kyotocabinet_src_dir}")

        # Set LD_PRELOAD environment variable if required
        env_vars = os.environ.copy()
        if preload_lib:
            env_vars["LD_PRELOAD"] = str(preload_lib)

        # Print the command being executed
        print(f"[RUN] {command}")

        # Make sure the command is executed within the correct directory
        result = subprocess.run(
            command,
            shell=True,
            check=True,
            stdout=subprocess.PIPE,  # Capturing output for display
            stderr=subprocess.PIPE,
            cwd=kyotocabinet_src_dir,  # Ensure we are in the correct directory
            env=env_vars,  # Pass the environment with LD_PRELOAD if necessary
        )

        # Output captured from benchmark run (stdout)
        print(result.stdout.decode())
        
    except subprocess.CalledProcessError as e:
        print(f"Error occurred while running the benchmark: {e}")
        raise



def main() -> None:
    """Main function for running Kyoto Cabinet benchmarks with Vsync locks."""
    platform = get_current_platform()
    build_ok = build_locks(platform=platform)

    # Define the path to the Kyoto Cabinet source directory
    kyotocabinet_src_dir = (get_curdir(__file__) / "deps/kyotocabinet-1.2.76/").resolve()

    # Standard benchmark command (without LD_PRELOAD)
    benchmark_cmd = "./benchmark_verify -t 63 -d 30"
    
    # Run the standard benchmark in the Kyoto Cabinet source directory (no LD_PRELOAD)
    print("Running standard benchmark...")
    run_benchmark(benchmark_cmd, kyotocabinet_src_dir)

    # Custom lock benchmark command (with LD_PRELOAD)
    custom_lock_cmd = f"LD_PRELOAD={build_ok}/libvmuslmutex.so {benchmark_cmd}"
    
    # Run the benchmark with the custom lock in the Kyoto Cabinet source directory
    print("Running benchmark with custom lock...")
    run_benchmark(custom_lock_cmd, kyotocabinet_src_dir, preload_lib=build_ok / "libvmuslmutex.so")


if __name__ == "__main__":
    main()